from leafinstance import load_leaf_instances
import os
import cv2
import numpy as np
import random
from PIL import Image

def apply_occlusion(image, mask, leaf, position):
    """将带透明通道的叶子实例覆盖到指定位置"""
    leaf_height, leaf_width = leaf.shape[:2]
    x, y = position
    
    # 确保叶子不会超出图像边界
    if x < 0:
        leaf_width += x
        x = 0
    if y < 0:
        leaf_height += y
        y = 0
    if x + leaf_width > image.shape[1]:
        leaf_width = image.shape[1] - x
    if y + leaf_height > image.shape[0]:
        leaf_height = image.shape[0] - y
    
    # 如果叶子完全在图像外，直接返回原图
    if leaf_width <= 0 or leaf_height <= 0:
        return image, mask
    
    leaf = leaf[:leaf_height, :leaf_width]
    
    # 分离alpha通道
    leaf_rgb = leaf[:, :, :3]
    leaf_alpha = leaf[:, :, 3] / 255.0
    
    # 使用alpha通道混合
    for c in range(3):
        image[y:y+leaf_height, x:x+leaf_width, c] = (
            image[y:y+leaf_height, x:x+leaf_width, c] * (1 - leaf_alpha) +
            leaf_rgb[:, :, c] * leaf_alpha
        )
    mask[y:y+leaf_height, x:x+leaf_width] = mask[y:y+leaf_height, x:x+leaf_width] * (1 - leaf_alpha)
    return image, mask

def random_rotate(leaf):
    """随机旋转叶子"""
    angle = random.uniform(0, 360)
    leaf = Image.fromarray(leaf)
    leaf = leaf.rotate(angle, expand=True)
    return np.array(leaf)

def random_resize(leaf):
    """随机缩放叶子"""
    scale = random.uniform(0.8, 1.2)  # 缩放范围在0.5到1.5倍之间
    h, w = leaf.shape[:2]
    new_size = (int(w * scale), int(h * scale))
    leaf = cv2.resize(leaf, new_size, interpolation=cv2.INTER_AREA)
    return leaf

if __name__ == '__main__':
    save_dir = r'D:\Jupyter\pytorch\d2l-zh\pytorch\appppppp\detection2\processed_blkbgleaf'
    image_dir = r'D:\Jupyter\pytorch\d2l-zh\pytorch\appppppp\detection2\train\images'
    mask_dir = r'D:\Jupyter\pytorch\d2l-zh\pytorch\appppppp\detection2\train\masks'
    
    # 读取带透明通道的叶子图像
    # loaded_instances = [cv2.imread(os.path.join(save_dir, fname), cv2.IMREAD_UNCHANGED) 
    #                    for fname in os.listdir(save_dir) if fname.endswith('.png')]
        # 读取三通道纯黑背景叶子图像
    loaded_instances = []
    for fname in os.listdir(save_dir):
        if fname.lower().endswith(('.png', '.jpg', '.jpeg')):
            leaf = cv2.imread(os.path.join(save_dir, fname), cv2.IMREAD_UNCHANGED)
            # 创建alpha通道，黑色背景为0，其他为255
            gray = cv2.cvtColor(leaf, cv2.COLOR_BGR2GRAY)
            _, alpha = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
            # 添加alpha通道
            leaf = cv2.merge((leaf[:, :, 0], leaf[:, :, 1], leaf[:, :, 2], alpha))
            loaded_instances.append(leaf)
    # random.seed(42)
    
    # 创建输出目录
    output_image_dir = os.path.join(os.path.dirname(image_dir), 'images2')
    output_mask_dir = os.path.join(os.path.dirname(mask_dir), 'masks2')
    os.makedirs(output_image_dir, exist_ok=True)
    os.makedirs(output_mask_dir, exist_ok=True)

    # 处理每张图像
    for image_name in [f for f in os.listdir(image_dir) 
                      if f.lower().endswith(('.png', '.jpg', '.jpeg')) 
                      and not f.endswith('_occul.png')]:
        image_path = os.path.join(image_dir, image_name)
        mask_path = os.path.join(mask_dir, image_name)
        print(image_name)
        image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        
        # 将mask转换为numpy数组并处理
        mask = np.array(mask)
        obj_ids = np.unique(mask)
        obj_ids = obj_ids[1:]  # 移除背景id
        
        # 如果没有检测到对象，跳过当前图像
        if len(obj_ids) == 0:
            print(f"警告：{image_name} 中没有检测到目标对象，跳过处理")
            continue
            
        # 确保masks是二维张量
        masks = mask == obj_ids[:, None, None]
        if masks.ndim == 1:
            masks = masks.reshape(-1, masks.shape[0])
        elif masks.size == 0:
            masks = np.zeros((0, mask.shape[0], mask.shape[1]), dtype=bool)
        num_leaves = random.randint(1, len(loaded_instances))
        # 随机选择指定数量的叶子（允许重复）
        selected_leaves = [random.choice(loaded_instances) for _ in range(num_leaves)]

        for leaf in selected_leaves:
            # 随机旋转叶子
            leaf = random_rotate(leaf)
            leaf = random_resize(leaf)
            
            # 获取叶子实例的有效区域
            leaf_mask = leaf[:, :, 3] > 0
            leaf_indices = np.where(leaf_mask)
            if len(leaf_indices[0]) > 0:
                # 计算叶子实例的边界
                leaf_min_y, leaf_max_y = np.min(leaf_indices[0]), np.max(leaf_indices[0])
                leaf_min_x, leaf_max_x = np.min(leaf_indices[1]), np.max(leaf_indices[1])
                leaf_height = leaf_max_y - leaf_min_y
                leaf_width = leaf_max_x - leaf_min_x
                
                # 在mask中随机选择一个苹果实例
                if len(obj_ids) > 0:
                    selected_obj_id = random.choice(obj_ids)
                    obj_mask = masks[obj_ids == selected_obj_id][0]
                    mask_indices = np.where(obj_mask > 0)

                    if len(mask_indices[0]) > 0:
                        # 计算实例的边界
                        min_y, max_y = np.min(mask_indices[0]), np.max(mask_indices[0])
                        min_x, max_x = np.min(mask_indices[1]), np.max(mask_indices[1])

                        # 计算叶子可以放置的位置范围
                        x_range = max_x - min_x - leaf_width
                        y_range = max_y - min_y - leaf_height

                        # 如果叶子实例太大，缩小到合适尺寸
                        if x_range < 0 or y_range < 0:
                            scale = min((max_x - min_x) / leaf_width, (max_y - min_y) / leaf_height)
                            new_size = (int(leaf.shape[1] * scale), int(leaf.shape[0] * scale))
                            leaf = cv2.resize(leaf, new_size, interpolation=cv2.INTER_AREA)
                            # 重新计算叶子实例边界
                            leaf_mask = leaf[:, :, 3] > 0
                            leaf_indices = np.where(leaf_mask)
                            leaf_min_y, leaf_max_y = np.min(leaf_indices[0]), np.max(leaf_indices[0])
                            leaf_min_x, leaf_max_x = np.min(leaf_indices[1]), np.max(leaf_indices[1])
                            leaf_height = leaf_max_y - leaf_min_y
                            leaf_width = leaf_max_x - leaf_min_x
                            x_range = max_x - min_x - leaf_width
                            y_range = max_y - min_y - leaf_height

                        # 在苹果边界内随机选择一个位置
                        if x_range > 0 and y_range > 0:
                            position = (
                                min_x + random.randint(0, x_range) - leaf_min_x,
                                min_y + random.randint(0, y_range) - leaf_min_y
                            )
                        else:
                            # 如果叶子实例仍然太大，放在中心位置
                            position = (
                                (min_x + max_x - leaf_width) // 2 - leaf_min_x,
                                (min_y + max_y - leaf_height) // 2 - leaf_min_y
                            )

                        # 应用遮挡
                        image, mask = apply_occlusion(image, mask, leaf, position)

                        # 检查是否完全覆盖了苹果实例
                        # 获取叶子覆盖区域的mask
                        y_start = max(position[1], 0)
                        y_end = min(position[1] + leaf_height, mask.shape[0])
                        x_start = max(position[0], 0)
                        x_end = min(position[0] + leaf_width, mask.shape[1])
                        
                        # 只检查叶子实际覆盖的区域（alpha > 0）
                        if y_end > y_start and x_end > x_start:
                            # 获取实际与图像重叠的叶子区域
                            leaf_region = leaf[:y_end-y_start, :x_end-x_start]
                            leaf_alpha = leaf_region[:, :, 3] > 0
                            covered_mask = mask[y_start:y_end, x_start:x_end]
                            # 只比较叶子实际覆盖的区域
                            if np.all(covered_mask[leaf_alpha] == 0):
                                print(f"警告：叶子完全覆盖了苹果实例，跳过该叶子")
                                continue

        # 保存增强后的图像和mask到新目录
        output_image_path = os.path.join(output_image_dir, image_name.replace('.png', '_occul.png'))
        output_mask_path = os.path.join(output_mask_dir, image_name.replace('.png', '_occul.png'))
        cv2.imwrite(output_image_path, image)
        cv2.imwrite(output_mask_path, mask)
